import os
import sys
import gc
import time
import torch
import argparse
from modelscope import snapshot_download
from diffsynth import ModelManager, FluxImagePipeline

def log(message):
    """简单的日志函数，带时间戳"""
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    print(f"[{timestamp}] {message}")

def free_memory():
    """释放内存的函数"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    log("内存已清理")

def download_model(model_name, model_base_path):
    """下载模型函数，带进度提示"""
    model_path = os.path.join(model_base_path, model_name)

    if os.path.exists(model_path):
        log(f"使用已下载的模型: {model_path}")
        return model_path

    log(f"开始下载模型到: {model_base_path}")
    log("下载可能需要一些时间，请耐心等待...")

    try:
        model_dir = snapshot_download(model_name, cache_dir=model_base_path)
        log(f"模型下载成功: {model_dir}")
        return model_dir
    except Exception as e:
        log(f"模型下载失败: {str(e)}")
        raise

def setup_pipeline(model_path, use_cpu=False):
    """设置和配置推理管道"""
    device = "cpu" if use_cpu else ("cuda" if torch.cuda.is_available() else "cpu")
    if device == "cpu":
        log("警告: 使用CPU进行推理，速度会很慢")
        # CPU模式下使用float32可能更稳定
        torch_dtype = torch.float32
    else:
        log(f"使用设备: {device}")
        # GPU模式使用float16以节省显存
        torch_dtype = torch.float16

    # 初始化模型管理器
    log("初始化模型管理器...")
    model_manager = ModelManager(torch_dtype=torch_dtype)

    # 加载模型
    log("加载模型中...")
    model_manager.load_models(
        file_path_list=[model_path],
        torch_dtype=torch_dtype,
        device=device
    )
    print(model_manager)
    # 创建pipeline
    log("创建推理管道...")
    pipe = FluxImagePipeline.from_model_manager(
        model_manager,
        device=device,
        torch_dtype=torch_dtype
    )

    if pipe is None:
        raise RuntimeError("创建推理管道失败")

    # 优化内存使用
    log("优化内存使用...")
    pipe.enable_cpu_offload()

    # 尝试启用量化
    if device != "cpu" and hasattr(pipe, 'dit') and pipe.dit is not None:
        log("启用模型量化以减少显存使用...")
        pipe.dit.quantize()

    return pipe

def generate_image(pipe, prompt, negative_prompt, output_path, steps=20, seed=None):
    """生成图像函数"""
    if seed is None:
        seed = int(time.time()) % 10000

    log(f"生成图像中... (步数: {steps}, 种子: {seed})")
    log(f"提示词: {prompt}")

    # 生成前清理内存
    free_memory()

    # 生成图像
    start_time = time.time()
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        seed=seed,
        num_inference_steps=steps,
        # scale=7.0  # 略微降低以减少资源使用
    )

    # 保存图像
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    image.save(output_path)

    elapsed = time.time() - start_time
    log(f"图像生成完成! 耗时: {elapsed:.2f}秒")
    log(f"图像已保存到: {output_path}")

    return image

def main():
    # 参数解析
    parser = argparse.ArgumentParser(description="简洁的FLUX模型图像生成工具")
    parser.add_argument("--prompt", type=str, default="一个美丽的女孩，穿着时尚的衣服，在阳光明媚的公园里",
                        help="生成图像的提示词")
    parser.add_argument("--negative_prompt", type=str,
                        default="nsfw, nude, bad quality, worst quality, low quality, blurry",
                        help="负面提示词")
    parser.add_argument("--steps", type=int, default=20,
                        help="推理步数，更少的步数速度更快但质量可能降低")
    parser.add_argument("--seed", type=int, default=None,
                        help="随机种子，不指定则使用随机值")
    parser.add_argument("--output", type=str, default="output/generated_image.jpg",
                        help="输出图像路径")
    parser.add_argument("--cpu", action="store_true",
                        help="强制使用CPU进行推理")
    parser.add_argument("--model_path", type=str,
                        default=os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                           "../../models"),
                        help="模型存储基础路径")
    args = parser.parse_args()

    try:
        # 设置模型路径
        model_base_path = args.model_path
        model_name = 'yiwanji/FLUX_xiao_hong_shu_ji_zhi_zhen_shi_V2'

        # 下载模型
        model_path = download_model(model_name, model_base_path)

        # 设置推理管道
        pipe = setup_pipeline(model_path, use_cpu=args.cpu)

        # 生成图像
        generate_image(
            pipe=pipe,
            prompt=args.prompt,
            negative_prompt=args.negative_prompt,
            output_path=args.output,
            steps=args.steps,
            seed=args.seed
        )

        # 清理资源
        del pipe
        free_memory()

        log("处理完成!")
        return 0

    except KeyboardInterrupt:
        log("用户中断处理")
        return 1
    except Exception as e:
        log(f"错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())